import os
import cv2
import numpy as np
import random


def getImgSize(path):
    img = cv2.imread(path)
    return img.shape


def getLabels(path):
    labels = []
    with open(path) as f:
        for line in f:
            labels.append(line.strip())
    labels = [decodeLabel(l) for l in labels if len(l) > 4]
    return labels


def decodeLabel(line):
    l = line.split()
    return [int(l[0]), float(l[1]), float(l[2]), float(l[3]), float(l[4])]


def createOutputImg(h, w, labels, path):
    img = np.zeros((h, w), dtype=np.uint8)
    for l in labels:
        if l[0] != 1 and l[0] != 2:
            continue
        x1, y1 = l[1] - l[3] / 2, l[2] - l[4] / 2
        x2, y2 = l[1] + l[3] / 2, l[2] + l[4] / 2
        x1, y1, x2, y2 = int(w * x1), int(h * y1), int(w * x2), int(h * y2)
        cv2.rectangle(img, (x1, y1), (x2, y2), 255, -1)
    cv2.imwrite(path, img)


if __name__ == "__main__":
    # 初始化路径
    root = "D:/pytorchPro/detection/dataset/"
    inpPath = root + "inputs/"
    outPath = root + "outputs/"
    labPath = root + "labels/"

    trainTxtPath = root + "train.txt"
    testTxtPath = root + "test.txt"

    trainRate = 0.8
    # 循环创建文件
    cnt = 0
    lis = os.listdir(inpPath)
    print("开始处理:", len(lis), ":", inpPath)
    for filename in lis:
        # 确定文件类型
        name, tp = filename.split(".")
        if tp != "jpg" and tp != "png" and tp != "jpeg":
            continue
        # 正式处理
        imgPath = inpPath + filename
        labelsPath = labPath + name + ".txt"
        h, w, _ = getImgSize(imgPath)
        labels = getLabels(labelsPath)
        createOutputImg(h, w, labels, outPath + name + ".jpg")
        # 分成训练集和测试集
        if random.random() < trainRate:
            with open(trainTxtPath, "a") as f:
                f.write(filename + "\n")
        else:
            with open(testTxtPath, "a") as f:
                f.write(filename + "\n")
        # 输出进度
        cnt += 1
        print(cnt, "in", len(lis), ":", filename)
